package com.example.shardingsphere.demo.sharding.strategy;

import com.example.shardingsphere.demo.client.OrderClient;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingValue;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;

/**
 * 复合分表算法
 */
public class OrderGoodsComplexShardingAlgorithm implements ComplexKeysShardingAlgorithm<String> {

    @Override
    public Collection<String> doSharding(Collection<String> tableNames, ComplexKeysShardingValue<String> shardingValue) {
        Collection<Integer> timeTypeValues = getShardingValue(shardingValue, "time_type");

        List<String> tables = new ArrayList<>();
        if (!timeTypeValues.isEmpty()) {
            for (Integer timeType : timeTypeValues) {
                if (timeType == 0) {
                    tables.add(shardingValue.getLogicTableName());
                } else {
                    getTables(tables, tableNames, shardingValue);
                }
            }
        } else {
            getTables(tables, tableNames, shardingValue);
        }
        return tables;
    }

    private void getTables(List<String> tables, Collection<String> tableNames, ComplexKeysShardingValue<String> shardingValue) {
        Collection<String> orderNoValues = getShardingValue(shardingValue, "order_no");

        if (!orderNoValues.isEmpty()) {
            for (String orderNo : orderNoValues) {
                if (OrderClient.isTodayOrder(orderNo)) {
                    tables.add(shardingValue.getLogicTableName());
                    continue;
                }
                tables.add(getTableName(tableNames, OrderClient.getTableIndex(orderNo)));
            }
        }
    }

    private String getTableName(Collection<String> tableNames, int tableIndex) {
        for (String tableName : tableNames) {
            if (tableName.endsWith(tableIndex + "")) {
                return tableName;
            }
        }
        throw new IllegalArgumentException();
    }

    private <R> Collection<R> getShardingValue(ComplexKeysShardingValue<String> shardingValues, final String key) {
        Collection<R> collection = (Collection<R>) shardingValues.getColumnNameAndShardingValuesMap().get(key);
        if (collection == null) {
            return Collections.emptyList();
        }
        return collection;
    }
}